from typing import List
import os
import tvm
from tvm import relax, tir
from tvm.ir import IRModule
from tzer import relax_seeds

import logging
import traceback

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def relax_to_tir(relax_mod, target="llvm") -> List[tir.PrimFunc]:
    if not isinstance(relax_mod, IRModule):
        logging.error(f"Expected IRModule, but got {type(relax_mod)}")
        raise TypeError(f"Expected IRModule, but got {type(relax_mod)}")

    tirs = []
    try:
        logging.info("Starting relax module transformation to tir functions.")
        with tvm.transform.PassContext(opt_level=0):
            logging.info("Applying LegalizeOps pass.")
            mod = relax.transform.LegalizeOps()(relax_mod)
            logging.info("Applying FuseOps pass.")
            mod = relax.transform.FuseOps()(mod)
            logging.info("Applying FuseTIR pass.")
            mod = relax.transform.FuseTIR()(mod)
            for gvar in mod.get_global_vars():
                if isinstance(mod[gvar], tir.PrimFunc):
                    tirs.append(mod[gvar])
        logging.info("Relax module transformation to tir functions completed successfully.")
    except Exception as e:
        logging.error(f"Error during relax module transformation to tir functions: {e}", exc_info=True)
        raise

    return tirs

def get_all_seeds() -> List[tir.PrimFunc]:
    tirs = []
    for i, item in enumerate(relax_seeds.MODEL_SEEDS):
        if isinstance(item, tuple):
            relax_mod = item[0]
        else:
            relax_mod = item
        if not isinstance(relax_mod, tvm.IRModule):
            print(f"[WARN] Skipping seed #{i}: not an IRModule (got {type(relax_mod)})")
            continue
        tirs.extend(relax_to_tir(relax_mod))
    return tirs

def get_lemon_seeds() -> List[tir.PrimFunc]:
    tirs = []
    lemon_seeds_dir = '/tzer/lemon_seeds'
    try:
        for tir_file in os.listdir(lemon_seeds_dir):
            tir_file_path = os.path.join(lemon_seeds_dir, tir_file)
            with open(tir_file_path, 'r') as f:
                data = f.read()
            mod = tvm.ir.load_json(data)
            tirs.extend([mod[v] for v in mod.get_global_vars()])
    except Exception as e:
        print(f"[ERROR] Something went wrong while loading Lemon seeds: {e}")
    return tirs